import json
from pathlib import Path
from openai import OpenAI, APIConnectionError, APIError, APIStatusError

class DialogueGenClient:
    def __init__(self, api_key, base_url=None):
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            max_retries=3  # 内置自动重试
        )

    def generate_dialogues(self, examples):
        """调用API生成对话"""
        system_msg = """你是一个对话生成专家，请严格按以下步骤工作：
1. 分析用户提供的对话示例中的语言风格、句式结构和应答模式
2. 生成20个全新场景的对话，保持相同风格但主题多样化
3. 严格使用原始JSON格式返回，每行一个包含input和output的对象"""
        
        user_content = "\n".join(
            [f"{ex['input']} -> {ex['output']}" for ex in examples]
        )

        try:
            response = self.client.chat.completions.create(
                model="ep-20250206141936-8l6xp",
                messages=[
                    {"role": "system", "content": system_msg},
                    {"role": "user", "content": user_content}
                ],
                response_format={"type": "json_object"},  # 强制JSON输出
                temperature=0.8,
                max_tokens=5000
            )
            return json.loads(response.choices[0].message.content)
        except (APIConnectionError, APIError, APIStatusError) as e:
            print(f"API错误: {e}")
            return None

def process_file(input_path, output_path, api_key, batch_size=20):
    """处理整个文件"""
    client = DialogueGenClient(api_key)
    output_path = Path(output_path)
    
    # 读取并分批
    with open(input_path, 'r', encoding='utf-8') as file:
        batch = []
        for line in file:
            try:
                data = json.loads(line)
                if validate_data(data):
                    batch.append(data)
                    if len(batch) >= batch_size:
                        process_batch(client, batch, output_path)
                        batch = []
            except json.JSONDecodeError:
                print(f"跳过无效行: {line.strip()}")
        
        if batch:  # 处理剩余数据
            process_batch(client, batch, output_path)

def validate_data(data):
    """验证数据格式"""
    return isinstance(data, dict) and \
           "input" in data and "output" in data and \
           isinstance(data["input"], str) and \
           isinstance(data["output"], str)

def process_batch(client, batch, output_path):
    """处理单个批次"""
    print(f"处理批次: 共{len(batch)}条样本...")
    result = client.generate_dialogues(batch)
    
    if not result or "dialogues" not in result:
        print("无效响应格式")
        return
    
    valid_count = 0
    with open(output_path, 'a') as f:
        for item in result["dialogues"]:
            if validate_data(item):
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
                valid_count += 1
    
    print(f"成功写入{valid_count}条数据")

if __name__ == "__main__":
    # 配置参数
    config = {
        "input_path": "xiaoyu.jsonl",
        "output_path": "generated.jsonl",
        "api_key": "xxxx"  # 从环境变量获取更安全
    }

    # 清空或创建输出文件
    Path(config["output_path"]).write_text('')
    
    process_file(**config)